package com.bruce.ai.service.impl;

import com.bruce.ai.service.ChatService;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.List;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Service
public class ChatServiceImpl implements ChatService {

    private static final Logger logger = LoggerFactory.getLogger(ChatServiceImpl.class);


    @Resource
    private ChatClient chatClient ;

    @Resource
    private VectorStore vectorStore ;

    @Value("classpath:/data/bikes.json")
    private org.springframework.core.io.Resource systemJsonResource ;

    @Value("classpath:/data/spring_ai_alibaba_quickstart.pdf")
    private org.springframework.core.io.Resource systemPdfResource ;

    @Override
    public boolean loadJsonData() {
        DocumentReader documentReader = new com.bruce.ai.reader.JsonReader(systemJsonResource) ;
        List<Document> documents = documentReader.get() ;
        if(documents == null || documents.isEmpty()){
            logger.error("load json data error");
            return false;
        }
        List<Document> splitDocuments = new TokenTextSplitter().apply(documents);
        vectorStore.add(splitDocuments);
        return true;
    }

    @Override
    public boolean loadPdfData() {
        DocumentReader documentReader = new PagePdfDocumentReader(systemPdfResource) ;
        List<Document> documents = documentReader.get() ;
        if(documents == null || documents.isEmpty()){
            logger.error("load pdf data error");
            return false;
        }
        List<Document> splitDocuments = new TokenTextSplitter().apply(documents);
        vectorStore.add(splitDocuments);
        return true;
    }

    @Override
    public Flux<String> generateStreamAsString(String message) {
        Flux<String> responseContent = chatClient.prompt()
                .system(s -> s.param("current_date", LocalDateTime.now().toString()))
                .advisors(a -> a.param(CHAT_MEMORY_RETRIEVE_SIZE_KEY,200))
                .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.query(message)))
                .user(message)
                .stream().content() ;

        return responseContent.concatWith(Flux.just("[complete]"));
    }
}
